from .args import args

heatmap_size = 64
boundary_num = 13

boundary_keys = ['chin', 'leb', 'reb', 'bon', 'breath', 'lue', 'lle', 'rue', 'rle', 'usul', 'lsul', 'usll', 'lsll']

interp_points_num = {
    'chin':   120,
    'leb':    32,
    'reb':    32,
    'bon':    32,
    'breath': 25,
    'lue':    25,
    'lle':    25,
    'rue':    25,
    'rle':    25,
    'usul':   32,
    'lsul':   32,
    'usll':   32,
    'lsll':   32
}

dataset_pdb_numbins = {
    '300W': 9,
    'AFLW': 17,
    'COFW': 7,
    'WFLW': 13
}

dataset_route = {
    '300W': args.dataset_route+'/300W/',
    'AFLW': args.dataset_route+'/AFLW/',
    'COFW': args.dataset_route+'/COFW/',
    'WFLW': args.dataset_route+'/WFLW/'
}

dataset_size = {
    '300W': {
        'train':            3148,
        'common_subset':    554,
        'challenge_subset': 135,
        'fullset':          689,
        '300W_testset':     600,
        'COFW68':           507  # 该数据集用于300W数据集上训练模型的测试
    },
    'AFLW': {
        'train':            20000,
        'test':             24386,
        'frontal':          1314
    },
    'COFW': {
        'train':            1345,
        'test':             507
    },
    'WFLW': {
        'train':            7500,
        'test':             2500,
        'pose':             326,
        'expression':       314,
        'illumination':     698,
        'makeup':           206,
        'occlusion':        736,
        'blur':             773
    }
}

kp_num = {
    '300W': 68, 
    'AFLW': 19,
    'COFW': 29, 
    'WFLW': 98
}

point_num_per_boundary = {
    '300W': [17., 5., 5., 4., 5., 4., 4., 4., 4., 7., 5., 5., 7.],
    'AFLW': [1.,  3., 3., 1., 2., 3., 3., 3., 3., 3., 3., 3., 3.],
    'COFW': [1.,  3., 3., 1., 3., 3., 3., 3., 3., 3., 1., 1., 3.],
    'WFLW': [33., 9., 9., 4., 5., 5., 5., 5., 5., 7., 5., 5., 7.]
}

boundary_special = {  # 有些边界线条使用的关键点和其他边界形成不连续交集，特殊处理
    'lle':  ['300W', 'COFW', 'WFLW'],
    'rle':  ['300W', 'COFW', 'WFLW'],
    'usll': ['300W', 'WFLW'],
    'lsll': ['300W', 'COFW', 'WFLW']
}

duplicate_point = {  # 需要重复使用的关键点的序号，从0开始计数
    '300W': {
        'lle':  36,
        'rle':  42,
        'usll': 60,
        'lsll': 48
    },
    'COFW': {
        'lle':  13,
        'rle':  17,
        'lsll': 21
    },
    'WFLW': {
        'lle':  60,
        'rle':  68,
        'usll': 88,
        'lsll': 76
    }
}

point_range = {  # notice: this is 'range', the later number pluses 1; the order is boundary order; index starts from 0
    '300W': [
        [0, 17],  [17, 22], [22, 27], [27, 31], [31, 36],
        [36, 40], [39, 42], [42, 46], [45, 48], [48, 55],
        [60, 65], [64, 68], [54, 60]
    ],
    'AFLW': [
        [0, 1],   [1, 4],   [4, 7],   [7, 8],   [8, 10],
        [10, 13], [10, 13], [13, 16], [13, 16], [16, 19],
        [16, 19], [16, 19], [16, 19]
    ],
    'COFW': [
        [0, 1],   [1, 4],   [5, 8],   [9, 10],  [10, 13],
        [13, 16], [15, 17], [17, 20], [19, 21], [21, 24],
        [25, 26], [26, 27], [23, 25]
    ],
    'WFLW': [
        [0, 33],  [33, 38], [42, 47], [51, 55], [55, 60],
        [60, 65], [64, 68], [68, 73], [72, 76], [76, 83],
        [88, 93], [92, 96], [82, 88]
    ]
}

flip_relation = {
    '300W': [
        [0, 16],  [1, 15],  [2, 14],  [3, 13],  [4, 12],  [5, 11],
        [6, 10],  [7, 9],   [8, 8],   [9, 7],   [10, 6],  [11, 5],
        [12, 4],  [13, 3],  [14, 2],  [15, 1],  [16, 0],  [17, 26],
        [18, 25], [19, 24], [20, 23], [21, 22], [22, 21], [23, 20],
        [24, 19], [25, 18], [26, 17], [27, 27], [28, 28], [29, 29],
        [30, 30], [31, 35], [32, 34], [33, 33], [34, 32], [35, 31],
        [36, 45], [37, 44], [38, 43], [39, 42], [40, 47], [41, 46],
        [42, 39], [43, 38], [44, 37], [45, 36], [46, 41], [47, 40],
        [48, 54], [49, 53], [50, 52], [51, 51], [52, 50], [53, 49],
        [54, 48], [55, 59], [56, 58], [57, 57], [58, 56], [59, 55],
        [60, 64], [61, 63], [62, 62], [63, 61], [64, 60], [65, 67],
        [66, 66], [67, 65]
    ],
    'AFLW': [
        [0, 0],   [1, 6],   [2, 5],   [3, 4],   [4, 3],   [5, 2],
        [6, 1],   [7, 7],   [8, 9],   [9, 8],   [10, 15], [11, 14],
        [12, 13], [13, 12], [14, 11], [15, 10], [16, 18], [17, 17],
        [18, 16]
    ],
    'COFW': [
        [0, 0],   [1, 7],   [2, 6],   [3, 5],   [4, 8],   [5, 3],
        [6, 2],   [7, 1],   [8, 4],   [9, 9],   [10, 12], [11, 11],
        [12, 10], [13, 19], [14, 18], [15, 17], [16, 20], [17, 15],
        [18, 14], [19, 13], [20, 16], [21, 23], [22, 22], [23, 21],
        [24, 24], [25, 25], [26, 26], [27, 28], [28, 27]
    ],
    'WFLW': [
        [0, 32],  [1, 31],  [2, 30],  [3, 29],  [4, 28],  [5, 27],
        [6, 26],  [7, 25],  [8, 24],  [9, 23],  [10, 22], [11, 21],
        [12, 20], [13, 19], [14, 18], [15, 17], [16, 16], [17, 15],
        [18, 14], [19, 13], [20, 12], [21, 11], [22, 10], [23, 9],
        [24, 8],  [25, 7],  [26, 6],  [27, 5],  [28, 4],  [29, 3],
        [30, 2],  [31, 1],  [32, 0],  [33, 46], [34, 45], [35, 44],
        [36, 43], [37, 42], [38, 50], [39, 49], [40, 48], [41, 47],
        [42, 37], [43, 36], [44, 35], [45, 34], [46, 33], [47, 41],
        [48, 40], [49, 39], [50, 38], [51, 51], [52, 52], [53, 53],
        [54, 54], [55, 59], [56, 58], [57, 57], [58, 56], [59, 55],
        [60, 72], [61, 71], [62, 70], [63, 69], [64, 68], [65, 75],
        [66, 74], [67, 73], [68, 64], [69, 63], [70, 62], [71, 61],
        [72, 60], [73, 67], [74, 66], [75, 65], [76, 82], [77, 81],
        [78, 80], [79, 79], [80, 78], [81, 77], [82, 76], [83, 87],
        [84, 86], [85, 85], [86, 84], [87, 83], [88, 92], [89, 91],
        [90, 90], [91, 89], [92, 88], [93, 95], [94, 94], [95, 93],
        [96, 97], [97, 96]
    ]
}

lo_eye_corner_index_x = {'300W': 72, 'AFLW': 20, 'COFW': 26, 'WFLW': 120}
lo_eye_corner_index_y = {'300W': 73, 'AFLW': 21, 'COFW': 27, 'WFLW': 121}
ro_eye_corner_index_x = {'300W': 90, 'AFLW': 30, 'COFW': 38, 'WFLW': 144}
ro_eye_corner_index_y = {'300W': 91, 'AFLW': 31, 'COFW': 39, 'WFLW': 145}
l_eye_center_index_x = {'300W': [72, 74, 76, 78, 80, 82], 'AFLW': 22, 'COFW': 54, 'WFLW': 192}
l_eye_center_index_y = {'300W': [73, 75, 77, 79, 81, 83], 'AFLW': 23, 'COFW': 55, 'WFLW': 193}
r_eye_center_index_x = {'300W': [84, 86, 88, 90, 92, 94], 'AFLW': 28, 'COFW': 56, 'WFLW': 194}
r_eye_center_index_y = {'300W': [85, 87, 89, 91, 93, 95], 'AFLW': 29, 'COFW': 57, 'WFLW': 195}
